from typing import Tuple,NamedTuple,Dict,Any,List,TypedDict,Optional
from ..judger import get_problem_path,StatusType,CompareType,ProblemType,get_problem_meta
import os,shutil
from pathlib import Path
import json
def prepare_workspace(worker_id:int):
    WORKSPACE=Path(__file__).parent.resolve()/"workspace"/f"ws_{worker_id}/"
    if WORKSPACE.exists():
        shutil.rmtree(WORKSPACE)
    WORKSPACE.mkdir()
    return WORKSPACE

def get_testcase_meta(problem_path,case_name):
    file:Path=problem_path/case_name/"meta.json"
    fp=file.open()
    return json.loads(fp.read())

GCC="g++"

class BodyMeta(TypedDict):
    code:str
    lang:str

class SubmissionMeta(TypedDict):
    body:BodyMeta
    problem_id:str
    type=ProblemType
    submission_id:int

class ProblemMeta(TypedDict):
    problem_id:str
    type:ProblemType
    data_count:int
    data:List[str]

class TestcaseMeta(TypedDict):
    problem_id:str
    name:str
    input:str
    answer:str
    compare:CompareType
    checker:Optional[str]
    points:int

class DetailMeta(TypedDict):
    status:StatusType
    score:int
    msg:str

class ResultMeta(TypedDict):
    status:StatusType
    score:int
    detail:Dict[str,DetailMeta]

class CompareResult(TypedDict):
    score:int
    msg:str

class Judge:
    workspace:Path=""
    problem:Path=""
    prob_meta:ProblemMeta
    case_meta:TestcaseMeta
    result:ResultMeta
    submit:SubmissionMeta

    LANGS={
        "C++14 O2":"CPP",
        "C++14"   :"CPP"
    }

    def answer(self):
        return (self.problem/self.case_meta["name"]/self.case_meta["answer"]).resolve()
    def input(self):
        return (self.problem/self.case_meta["name"]/self.case_meta["input"]).resolve()
    def checker(self):
        return (self.problem/self.case_meta["name"]/self.case_meta["checker"]).resolve()

    # Use system command
    def syscmd_cmp(self,output_f:Path)->CompareResult:
        import os
        if os.name=="nt":
            f=os.system(f"fc {self.answer()} {output_f.resolve()} > nul")
        elif os.name=="posix":
            f=os.system(f"diff {self.answer()} {output_f.resolve()} > /dev/null")
        else:
            raise NotImplementedError("Not supported OS")
        result:CompareResult={
            'score':(self.case_meta["points"] if f==0 else 0),
            'msg':"No furtuer infomation."
        }
        return result
    
    # Checker / <input> <output> <answer> / print to stdout
    CHECKER_TL=5
    #don't know if it can work with testlib
    def checker_cmp(self,output_f:Path)->CompareResult:
        import subprocess
        res=subprocess.run(
            [self.checker(),self.input(),output_f.resolve(),self.answer()],
            capture_output=True,
            timeout=self.CHECKER_TL,
            text=True
        )
        exit_code = res.returncode
        stdout = res.stdout.strip()
        
        if '|' in stdout:
            score_part,message=stdout.split('|', 1)
        else:
            score_part,message=stdout,"No furtuer info"
        
        result:CompareResult={
            'score':0,
            'msg':message
        }
        if exit_code==0:
            if score_part.replace('.', '').isdigit():  # 是分数
                result["score"]=int(float(score_part)*int(self.case_meta["points"]))
            else:
                result["score"]=self.case_meta["points"]
        return result

    def compare(self,output_f:Path)->CompareResult:
        if self.case_meta["compare"]==CompareType.SYSCMD:
            return self.syscmd_cmp(output_f)
        if self.case_meta["compare"]==CompareType.CHECKER:
            return self.checker_cmp(output_f)

    #cpp
    CPP_FLAGS={
        "C++14 O2":"-DONLINE_JUDGE -std=c++14 -O2",
        "C++14"   :"-DONLINE_JUDGE -std=c++14"
    }
    def init_cpp(self):
        self.code_file=(self.workspace/"code.cpp").resolve()
        with self.code_file.open("w+") as file_c:
            file_c.write(self.submit["body"]["code"])
            file_c.close()
        lang=self.submit["body"]["lang"]
        if lang not in self.CPP_FLAGS:
            raise KeyError("Invalid lang")
        self.exe_file=(self.workspace/"exe.exe").resolve()
        import os
        os.system(f"{GCC} {self.CPP_FLAGS[lang]} -o {self.exe_file} {self.code_file}")

    def work_cpp(self):
        result:DetailMeta={
            'status':StatusType.FAILED,
            'score':0,
            'msg':'NO INFO'
        }
        import subprocess
        # not completed
        

    # subtasks and grading script is not supported currently
    def make_result(self):
        self.result['status']="AC"
        self.result['score']=0
        for i in self.prob_meta["data"]:
            if i not in self.result['detail']:
                self.result["status"]=StatusType.UNACCEPTED
            else:
                if self.result['detail'][i]["status"]!=StatusType.ACCEPTED:
                    self.result["status"]=StatusType.UNACCEPTED
                self.result['score']+=self.result['detail'][i]['score']
    
    def judge(self,worker_id:int,data:SubmissionMeta):
        print(f"Judging {data["submission_id"]} problem {data['problem']}")
        
        #Prepare workspace
        self.workspace=prepare_workspace(worker_id)
        self.problem=get_problem_path(data["problem"])
        self.prob_meta:ProblemMeta=get_problem_meta(data["problem"])
        self.submit=data

        self.result:ResultMeta={
            "status":StatusType.FAILED,
            "score":0,
            "detail":{}
        }

        L=self.LANGS[data["body"]["lang"]]
        if L=="CPP":
            self.init_cpp()
        
        for i in self.prob_meta["data"]:
            self.case_meta=get_testcase_meta(self.problem,i)
            if L=="CPP":
                self.cppwork(i)
        
        self.make_result()
        return self.result

from multiprocessing import Queue,Process,Pool,Manager
import sys
def worker(worker_id:int,current_work:Dict[int,int],pending:Queue,result:Queue,task_data:Dict[int,Any]):
    print(f"Working #{worker_id} started.")
    sys.stdout.flush()
    if worker_id in current_work:
        del current_work[worker_id]
    while True:
        try:
            rid=pending.get()
        except:
            print("Quiting")
            return
        current_work[worker_id]=rid
        try:
            res=judge(worker_id,task_data[rid])
            res["submission_id"]=rid
            result.put(res)
        except:
            print(f"Exception caught during handling #{rid}")
            result.put({"status":'SE',"score":0,"submission_id":rid,"detail":{}})
        finally:
            del task_data[rid]
            del current_work[worker_id]